import tensorflow as tf
from keras_preprocessing.image import ImageDataGenerator
from keras import layers, Model, Sequential, activations, optimizers, losses
from keras.datasets.fashion_mnist import load_data

# 加载Fashion-MNIST数据集
(x_train, y_train), (x_test, y_test) =load_data()

# 将数据reshape为适合模型输入的形状并标准化到[0, 1]范围
x_train = x_train.reshape(-1, 28, 28, 1) / 255.0
x_test = x_test.reshape(-1, 28, 28, 1) / 255.0

# 定义数据增强
data_augmentation = ImageDataGenerator(
    width_shift_range=0.1,
    height_shift_range=0.1,
    horizontal_flip=True,
)

# 定义AlexNet模型
class AlexNet(Model):
    def __init__(self):
        super(AlexNet, self).__init__()
        self.conv = Sequential([
            layers.Conv2D(filters=6, kernel_size=(3, 3), activation='relu', padding='same'),
            layers.MaxPooling2D(),
            layers.Conv2D(filters=16, kernel_size=(3, 3), activation='relu', padding='same'),
            layers.MaxPooling2D(),
            layers.Conv2D(filters=24, kernel_size=(3, 3), activation='relu', padding='same'),
            layers.Conv2D(filters=24, kernel_size=(3, 3), activation='relu', padding='same'),
            layers.Conv2D(filters=16, kernel_size=(3, 3), activation='relu', padding='same'),
            layers.MaxPooling2D()
        ])
        self.flat = layers.Flatten()
        self.fc = Sequential([
            layers.Dense(128, activation='relu'),
            layers.Dense(64, activation='relu'),
            layers.Dense(64, activation='relu'),
            layers.Dense(10, activation='softmax')
        ])

    def call(self, inputs):
        x = self.conv(inputs)
        x = self.flat(x)
        x = self.fc(x)
        return x

# 实例化和构建模型
model = AlexNet()
model.build(input_shape=(None, 28, 28, 1))  # 输入形状为28x28x1
model.summary()

# 编译模型
model.compile(optimizer=optimizers.Adam(),
              loss=losses.sparse_categorical_crossentropy,
              metrics=['acc'])

# 训练模型
model.fit(data_augmentation.flow(x_train, y_train, batch_size=64), epochs=10)

# 在测试集上评估模型
test_loss, test_acc = model.evaluate(x_test, y_test)
print(f'Test accuracy: {test_acc}')

# 保存模型权重
model.save_weights('alexnet_fashion_mnist.h5')
